# 导入 pymilvus 集合模式、字段模式、数据类型、集合和连接模块
from pymilvus import CollectionSchema, FieldSchema, DataType, Collection, connections

# 导入随机数生成模块
import random

# 连接到 Milvus 数据库，指定主机、端口和数据库名称
connections.connect("default", host="42.194.164.248", port="19530", db_name="rensheng")

# 定义包含标量字段的 schema
# 创建字段列表，包含主键ID、向量字段和分类字段
fields = [
    # 定义主键字段，使用64位整数类型，自动生成ID
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    # 定义向量字段，使用128维浮点向量类型
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=128),
    # 定义分类字段，使用64位整数类型
    FieldSchema(name="category", dtype=DataType.INT64),
]

# 创建集合模式对象
schema = CollectionSchema(fields=fields, description="Mixed collection schema with scalar field")

collection = Collection("mixed_collection", schema)

# 插入带标量字段的数据
# 生成100个128维随机向量
vectors = [[random.random() for _ in range(128)] for _ in range(100)]
# 生成100个1到5之间的随机分类标签
categories = [random.randint(1, 5) for _ in range(100)]
# 将向量和分类数据组织成插入格式
data = [vectors, categories]
# 将数据插入到集合中
collection.insert(data)

# 创建索引和加载
# 定义索引参数：使用L2距离度量，IVF_FLAT索引类型，聚类数量为64
index_params = {"metric_type": "L2", "index_type": "IVF_FLAT", "params": {"nlist": 64}}
# 为向量字段创建索引
collection.create_index("embedding", index_params)
# 将集合加载到内存中以支持搜索操作
collection.load()

# 带过滤条件的搜索
# 生成一个128维随机查询向量
query_vector = [random.random() for _ in range(128)]
# 定义搜索参数：使用L2距离度量，nprobe参数为10
search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
# 执行带过滤条件的向量搜索，只搜索分类为2、3、4的数据
results = collection.search(
    # 查询向量数据
    data=[query_vector],
    # 指定搜索的向量字段名
    anns_field="embedding",
    # 搜索参数配置
    param=search_params,
    # 返回前5个最相似的结果
    limit=5,
    # 过滤条件：只搜索分类在[2,3,4]范围内的数据
    expr="category in [2, 3, 4]",
)
# 遍历搜索结果
for idx,hits in enumerate(results):
    print(f"Search result for query vector {idx + 1}:")
    # 遍历每个命中结果
    for i, hit in enumerate(hits):
        print(f"Top {i + 1} hit:")
        # 打印结果ID和距离值
        print(f"id: {hit.id}, distance: {hit.distance}")